# %%
import numpy as np
import scipy.io as sio
import os
import time
from variables import (ku_grid, ky_grid, kv_grid, Nu, Ny, Nv, m0, n0, l0, version_count,
                       save_address, num_of_band, V0, Vy_List, Omega0, mz_list, psi_list, AtomTemperture,
                       AtomDensityHot)
import HotAtomDist_3D
import WeylPoint_Count
import Diag_Hamiltonian_No_psi
os.system('chcp 65001')
# %% Initial Settings
start_time = time.time()
start_time_struct = time.localtime(start_time)
for Vy_Index in range(len(Vy_List)):
    Vy = Vy_List[Vy_Index]
    for del_psi_Index in range(len(psi_list)):
        del_psi = psi_list[del_psi_Index]
        Vxz = V0 * np.cos(del_psi)
        Omega_xy = Omega0 * np.sin(del_psi / 2 + np.pi / 4)
        Omega_zy = Omega0 * np.cos(del_psi / 2 + np.pi / 4)
        t1 = Vxz
        t2 = Vy
        t3 = Omega_xy
        t4 = Omega_zy
        for mz_Index in range(len(mz_list)):
            mz = mz_list[mz_Index]
            ns = np.zeros((num_of_band, Nu, Ny, Nv), dtype=np.float64)
            parameters = str(m0) + 'x'+str(n0) + 'x' + str(l0) + ';' + str(num_of_band) + 'x'+str(Nu) + 'x' + \
                str(Ny) + 'x' + str(Nv) + ';V_xz=' + '%.3f' % t1 + 'E_r,V_y=' + '%.3f' % t2 + 'E_r,Oxy=' + \
                '%.3f' % t3 + 'E_r,Ozy=' + '%.3f' % t4 + \
                'E_r,T=%.3eK,n=%.3em-3' % (AtomTemperture, AtomDensityHot) + \
                'del_psi=%.3fpi,' % (del_psi / np.pi) + 'mz=%.3fEr' % mz
            print('calculating: mz=%.3f, del_psi=%.3fpi, Vy=%.3f' % (mz, del_psi / np.pi, Vy))

            if 0 and os.path.exists(save_address + parameters + 'no_states_v' + str(version_count) + '.npz') and \
                    os.path.exists(save_address + parameters + 'no_states_v' + str(version_count) + '.mat'):
                continue
            energies, ESigmaZ = Diag_Hamiltonian_No_psi.diag_no_states(mz, num_of_band, t1, t2, t3, t4)
            gap = energies[1, ...] - energies[0, ...]
            if Nu > 10 and Ny > 10 and Nv > 10:
                (mu, ns, nsp) = HotAtomDist_3D.ChemicalPotential(energies, AtomTemperture, AtomDensityHot)
                Spin_Texture = np.sum(nsp * ESigmaZ, 0)
            else:
                (mu, ns, nsp, Spin_Texture) = (0, 0, 0, 0)
            Weyl_Points = []
            Number_of_WeylPoints = 0
            for ku_index in range(len(ku_grid)):
                if ku_grid[ku_index] == 1:
                    continue
                gap2d = gap[ku_index, ...]
                Weyl_Points_ku, Number_of_WeylPoints_ku = WeylPoint_Count.Count(gap2d, ky_grid, kv_grid, 0.01)
                Weyl_Points_kup = []
                if Number_of_WeylPoints_ku != 0:
                    for point in Weyl_Points_ku:
                        Weyl_Points_kup.append((ku_grid[ku_index], point[0], point[1]))
                Number_of_WeylPoints += Number_of_WeylPoints_ku
                Weyl_Points.extend(Weyl_Points_kup)
            np.savez(save_address + parameters + 'no_states_v' + str(version_count) + '.npz',
                     ku_grid=ku_grid, ky_grid=ky_grid, kv_grid=kv_grid,
                     energies=energies, ESigmaZ=ESigmaZ, ns=ns, Spin_Texture=Spin_Texture,
                     Weyl_Points=Weyl_Points, Number_of_WeylPoints=Number_of_WeylPoints,
                     parameters=parameters, mz=mz)
            sio.savemat(save_address + parameters + 'no_states_v' + str(version_count) + '.mat',
                        {'ku_grid': ku_grid, 'ky_grid': ky_grid, 'kv_grid': kv_grid,
                         'energies': energies, 'ESigmaZ': ESigmaZ, 'ns': ns, 'nsp': nsp, 'Spin_Texture': Spin_Texture,
                         'parameters': parameters, 'mz': mz})
end_time = time.time()
end_time_struct = time.localtime(end_time)
print("Start time: ", start_time)
print("End time: ", end_time)
print(version_count)
